package com.lixw.langchain.controller;

import cn.hutool.core.date.DateUtil;
import com.lixw.langchain.entity.LawPrompt;
import com.lixw.langchain.service.LawAssistant;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.Map;

@RestController
@Slf4j
public class ChatPromptController
{
    @Resource
    private LawAssistant lawAssistant;
    @Resource
    private ChatModel chatModel;

    // http://localhost:9009/chatprompt/test1
    @GetMapping(value = "/chatprompt/test1")
    public String test1()
    {
        String chat = lawAssistant.chat("什么是知识产权？",2000);
        System.out.println(chat);

        String chat2 = lawAssistant.chat("什么是java？",2000);
        System.out.println(chat2);

        String chat3 = lawAssistant.chat("介绍下西瓜和芒果",2000);
        System.out.println(chat3);

        String chat4 = lawAssistant.chat("飞机发动机原理",2000);
        System.out.println(chat4);

        return "success : "+ DateUtil.now()+"<br> \n\n chat: "+chat+"<br> \n\n chat2: "+chat2;
    }

    @GetMapping("/chatprompt/test2")
    public String test2(){
        LawPrompt lawPrompt = new LawPrompt();
        lawPrompt.setLegal("知识产权");
        lawPrompt.setQuestion("TRIP协议是什么");

        String chat = lawAssistant.chat(lawPrompt);

        log.info(chat);

        return "success : "+ DateUtil.now()+"<br> \n\n chat: "+chat;
    }

    @GetMapping(value = "/chatpropmt/test3")
    public String test3(){
        String role = "外科医生";
        String question = "牙疼";

        PromptTemplate template = PromptTemplate.from("你是一个{{it}}助手，{{question}}怎么办？");

        Prompt prompt = template.apply(Map.of("it", role, "question", question));

        UserMessage userMessage = prompt.toUserMessage();

        ChatResponse chatResponse = chatModel.chat(userMessage);

        System.out.println(chatResponse.aiMessage().text());
        //4.2 前台返回
        return "success : "+ DateUtil.now()+"<br> \n\n chat: "+chatResponse.aiMessage().text();
    }

    @GetMapping(value = "/chatpropmt/test4")
    public String test4(){
        String question = "1 + 1等于几?";

        PromptTemplate template = PromptTemplate.from("你是一个小学老师助手，{{it}}怎么办？");

        Prompt prompt = template.apply(question);

        UserMessage userMessage = prompt.toUserMessage();

        ChatResponse chatResponse = chatModel.chat(userMessage);

        System.out.println(chatResponse.aiMessage().text());
        //4.2 前台返回
        return "success : "+ DateUtil.now()+"<br> \n\n chat: "+chatResponse.aiMessage().text();
    }

}